from typing import Any, Coroutine
from arq import create_pool
from arq.connections import ArqRedis, RedisSettings
import arq.worker
from axiom_boot.conf.manager import settings
from axiom_boot.di.decorators import service
from axiom_boot.di.application_context import ApplicationContext
from axiom_boot.logging.setup import get_logger

logger = get_logger(__name__)


@service()
class ArqManager:
    _pool: ArqRedis | None = None
    task_functions: list = []
    app_context: ApplicationContext | None = None  # 新增: 用于暂存 DI 容器

    def __init__(self):
        pass

    async def _get_pool(self) -> ArqRedis:
        if self._pool is None:
            redis_url = settings.task.redis_url or settings.redis_url
            if not redis_url:
                raise ValueError("未配置 ARQ 的 Redis URL (task.redis_url 或 redis_url 必须提供一个)。")
            
            redis_settings = RedisSettings.from_dsn(str(redis_url))
            logger.info(f"正在为 ARQ 创建 Redis 连接池，地址: {redis_settings.host}:{redis_settings.port}")
            self._pool = await create_pool(redis_settings)
        return self._pool

    async def enqueue_task(self, function_name: str, *args: Any, **kwargs: Any) -> Coroutine[Any, Any, Any] | None:
        """
        将一个任务添加到队列中。

        Args:
            function_name: 任务函数的名称 (例如 'my_task_function')。
            *args: 任务函数的位置参数。
            **kwargs: 任务函数的关键字参数。
        
        Returns:
            一个 Job 实例，如果成功入队。
        """
        pool = await self._get_pool()
        return await pool.enqueue_job(function_name, *args, **kwargs)

    async def run_worker(self, app_context: ApplicationContext, queue_name: str | None = None):
        """
        启动 ARQ Worker 监听任务。

        Args:
            app_context: Axiom Boot 的应用上下文实例 (DI 容器)。
            queue_name: (可选) 要监听的特定队列名称。
        """
        self.app_context = app_context  # [关键修复] 将 DI 容器保存在实例上，待 startup 时使用

        worker_settings = {
            "functions": self.task_functions,
            "on_startup": self._on_worker_startup,
            "on_shutdown": self._on_worker_shutdown,
            "redis_pool": await self._get_pool(),
            # [诊断性修复] 临时禁用健康检查，以防止 Playwright 等重操作导致任务被错误地取消
            "health_check_interval": 0,
        }
        if queue_name:
            worker_settings["queue_name"] = queue_name

        worker = arq.worker.Worker(**worker_settings)
        await worker.main()

    async def _on_worker_startup(self, ctx: dict):
        """
        在 worker 启动时被调用，负责将 DI 容器注入到 worker 上下文中。
        """
        logger.info("ARQ Worker 启动...")
        if self.app_context:
            ctx['axiom_app_context'] = self.app_context  # [关键修复] 执行注入
            logger.info("Axiom DI 容器已成功注入 Worker 上下文。")
        else:
            logger.error("未能找到 app_context 来注入 Worker 上下文！")

    async def _on_worker_shutdown(self, ctx: dict):
        logger.info("ARQ Worker 关闭...")
        await self.close()

    async def close(self):
        """关闭 Redis 连接池。"""
        if self._pool:
            await self._pool.close()
            self._pool = None
            logger.info("ARQ Redis 连接池已关闭。") 